#include "mm.h"
#include "sched.h"
#include "common.h"

page_frame_t page_frame[PAGE_FRAME_NUM];
uint64_t wrong_pid;

void init_page_table(pcb_t *pcb)
{
    int i;
    uint64_t pfn = 0x20000;
    for (i = 0; i < PAGE_MAX_NUM; i++)
    {
        // only used in TASK 1
        // pt[i].pt_entry = (pfn << 7) | (PTE_C << 4) | (PTE_D << 3) | (PTE_V << 2);
        // pfn++;
        pcb->page_table[i].vaddr = i * 0x1000;
        pcb->page_table[i].pinned = 0;
    }
    // pinned stack space into physical prame
    pcb->page_table[0xf].pinned = 1;
    pcb->page_table[0xe].pinned = 1;
}

void tlb_modify_helper()
{
    // we will not hit this exception so simply print it
    printk("The TLB modification exception is triggered\n");
}

void init_TLB(void)
{
    // fill TLB statically
    int i;
    uint64_t vpn2 = 0x0;
    uint64_t asid = 0, pagemask = 0;
    uint32_t index = 32;
    uint64_t pfn = 0x20000; // from 0.5G physical address
    // !!!Important!!!
    // C core: We can only init 32 ~ 63 TLB, because 0 ~ 31 TLB are
    // the terms that the system maps to the IO space and the PCIE space
    for (i = 32; i < TLB_NUM; i++, vpn2++, index++, pfn++)
    {
        // set_cp0_entryhi(vpn2 << 13 | asid & 0xff);
        // set_cp0_entrylo0(pfn << 6 | PTE_C << 3 | PTE_D << 2 | PTE_V << 1 | PTE_G);
        // pfn++;
        // set_cp0_entrylo1(pfn << 6 | PTE_C << 3 | PTE_D << 2 | PTE_V << 1 | PTE_G);
        set_cp0_entryhi(ENTRYHI_INVALID);
        set_cp0_entrylo0(0);
        set_cp0_entrylo1(0);

        set_cp0_pagemask(pagemask);
        set_cp0_index(index);
        asm volatile("tlbwi");
    }
}

void physical_frame_initial(uint64_t start_addr, uint64_t end_addr)
{
    uint64_t addr = start_addr;
    uint64_t sd_addr = 0x200;
    int i;
    clock_ptr = NULL;
    for (i = 0; addr + PAGE_SIZE <= end_addr; i++, addr += PAGE_SIZE)
    {
        page_frame[i].pfn = i;
        page_frame[i].paddr = addr;
        page_frame[i].sd_card = 0;
        page_frame[i].vpte = NULL;
        page_frame[i].key = 0;
        page_frame[i].asid = 0;
        page_frame_init(&page_frame[i]);
    }
    // init page frame in SD card
    for (; sd_addr + PAGE_SIZE <= SD_MAX_ADDR; i++, sd_addr += PAGE_SIZE)
    {
        page_frame[i].pfn = i;
        page_frame[i].paddr = sd_addr;
        page_frame[i].sd_card = 1;
        page_frame[i].vpte = NULL;
        page_frame[i].key = 0;
        page_frame[i].asid = 0;
        page_frame_init(&page_frame[i]);
    }
}

page_frame_t *swap_page_frame()
{
    page_frame_t *choose;
    clock_ptr = (page_frame_t *)busy_list.head;
    while (1)
    {
        // get one page_frame in memory which is busy
        if (clock_ptr->sd_card == 0)
        {
            if (clock_ptr->vpte != NULL)
            {
                if (clock_ptr->vpte->pinned == 0)
                {
                    choose = clock_ptr;
                    if (clock_ptr != busy_list.tail)
                    {
                        clock_ptr = clock_ptr->next;
                    }
                    else
                    {
                        clock_ptr = busy_list.head;
                    }
                    return choose;
                }
                else
                {
                    if (clock_ptr != busy_list.tail)
                    {
                        clock_ptr = clock_ptr->next;
                    }
                    else
                    {
                        clock_ptr = busy_list.head;
                    }
                }
            }
            else
            {
                choose = clock_ptr;
                if (clock_ptr != busy_list.tail)
                {
                    clock_ptr = clock_ptr->next;
                }
                else
                {
                    clock_ptr = busy_list.head;
                }
                return choose;
            }
        }
        else
        {
            if (clock_ptr != busy_list.tail)
            {
                clock_ptr = clock_ptr->next;
            }
            else
            {
                clock_ptr = busy_list.head;
            }
        }
    }
    // one problem: if every page frame in busy_list is in SD card, the program will loop infinitely...
    // yeah, but we programmer will avoid this situation... ^_^
}

void do_tlb_refill()
{
    uint64_t entrylo0, entrylo1;
    uint64_t entryhi;
    uint32_t index;
    uint64_t context = get_cp0_context();
    uint64_t vpn2, asid;
    context = (context & 0x7fffff) << 9;
    entryhi = (context & 0xffffffffffffff00) | (current_running[core_id]->pid & 0xff);

    // if S = 1 and V = 1(which means that there is a valid mapping but physical page frame is in SD card), then swap
    if (((current_running[core_id]->page_table[context >> 12].pt_entry & 0x1) == 1 && (current_running[core_id]->page_table[context >> 12].pt_entry & 0x4) == 0x4) ||
        ((current_running[core_id]->page_table[context >> 12 | 1].pt_entry & 0x1) == 1 && (current_running[core_id]->page_table[context >> 12 | 1].pt_entry & 0x4) == 0x4))
    {
        PTE_t *pte_0 = &(current_running[core_id]->page_table[context >> 12]);
        PTE_t *pte_1 = &(current_running[core_id]->page_table[context >> 12 | 1]);
        page_frame_t *frame0 = &page_frame[pte_0->f_pfn];
        page_frame_t *frame1 = &page_frame[pte_1->f_pfn];

        page_frame_t *pframe0;
        page_frame_t *pframe1;

        // choose a free page frame from memory
        page_frame_t *item0 = (page_frame_t *)free_list.head;
        while (item0 != NULL)
        {
            if (item0->sd_card == 0)
            {
                break;
            }
            item0 = (page_frame_t *)item0->next;
        }
        if (item0 != NULL)
        {
            pframe0 = item0;
        }
        else // if no free page in memory, then swap a page_frame in memory (CLOCK)
        {
            pframe0 = swap_page_frame();
        }

        // find the adjacent page table of pframe0->vpte, also swap its physical page frame out
        PTE_t *adj_pte;
        if ((pframe0->vpte->vaddr & ODD_PAGE) == ODD_PAGE)
        {
            adj_pte = &(current_running[core_id]->page_table[(pframe0->vpte->vaddr >> 12) - 1]);
        }
        else
        {
            adj_pte = &(current_running[core_id]->page_table[(pframe0->vpte->vaddr >> 12) + 1]);
        }
        pframe1 = &(page_frame[adj_pte->f_pfn]);

        frame_exchange(pframe0, frame0, pframe1, frame1);

        frame0->vpte->pt_entry &= 0xfffffffe;
        frame1->vpte->pt_entry &= 0xfffffffe;
    }
    // write TLB
    entrylo0 = (current_running[core_id]->page_table[context >> 12].pt_entry) >> 1;
    entrylo1 = (current_running[core_id]->page_table[context >> 12 | 1].pt_entry) >> 1;
    index = rand_index();
    set_cp0_index(index);
    set_cp0_entryhi(entryhi);
    set_cp0_entrylo0(entrylo0);
    set_cp0_entrylo1(entrylo1);
    asm volatile("tlbwi");
}

void do_tlb_invalid()
{
    uint64_t entrylo0, entrylo1;
    uint64_t entryhi;
    uint64_t index = get_cp0_index();
    uint64_t context = get_cp0_context();
    context = (context & 0x7fffff) << 9;

    if (!(current_running[core_id]->page_table[context >> 12].pt_entry & (PTE_V << 2)))
    {
        //printk("\ralloc first!\n");
        if (free_list.head != NULL)
        {
            page_frame_t *frame0;
            frame0 = page_alloc();

            // rewrite page_table (even)
            current_running[core_id]->page_table[context >> 12].pt_entry = ((frame0->paddr & 0xfffff000) >> 5 | (PTE_C << 4) | (PTE_D << 3) | (PTE_V << 2) | (frame0->sd_card));
            current_running[core_id]->page_table[context >> 12].f_pfn = frame0->pfn;
            frame0->vpte = &current_running[core_id]->page_table[context >> 12];
            frame0->asid = current_running[core_id]->pid;
        }
    }
    else
    {
        //printk("no hit 0\n");
    }
    if (!(current_running[core_id]->page_table[context >> 12 | 1].pt_entry & (PTE_V << 2)))
    {
        //printk("\ralloc second!\n");
        if (free_list.head != NULL)
        {
            page_frame_t *frame1;
            frame1 = page_alloc();

            // rewrite page_table (odd)
            current_running[core_id]->page_table[context >> 12 | 1].pt_entry = ((frame1->paddr & 0xfffff000) >> 5 | (PTE_C << 4) | (PTE_D << 3) | (PTE_V << 2) | (frame1->sd_card));
            current_running[core_id]->page_table[context >> 12 | 1].f_pfn = frame1->pfn;
            frame1->vpte = &current_running[core_id]->page_table[context >> 12 | 1];
            frame1->asid = current_running[core_id]->pid;
        }
    }
    else
    {
        //printk("no hit 1\n");
    }
    // rewrite TLB
    entryhi = (context & 0xffffff00) | (current_running[core_id]->pid & 0xff);
    entrylo0 = (current_running[core_id]->page_table[context >> 12].pt_entry) >> 1;
    entrylo1 = (current_running[core_id]->page_table[context >> 12 | 1].pt_entry) >> 1;
    set_cp0_entryhi(entryhi);
    set_cp0_entrylo0(entrylo0);
    set_cp0_entrylo1(entrylo1);
    asm volatile("tlbp");
    index = get_cp0_index();
    if (!(index & 0x80000000))
    {
        index = index;
    }
    else
    {
        index = rand_index();
    }
    set_cp0_index(index);
    asm volatile("tlbwi");
}

// The core function!!! The entry of TLB exception!!!
void tlb_exception_helper()
{
    uint64_t entryhi;
    uint32_t index;
    uint64_t pid = current_running[core_id]->pid;

    // FIRST STEP: get index
    uint64_t context = get_cp0_context();
    context = (context & 0x7fffff) << 9;
    entryhi = (context & 0xffffffffffffff00) | (pid & 0xff);

    set_cp0_entryhi(entryhi); // bad vaddr move from [22:4] to [31:13], with origin ASID on last 8 bits
    asm volatile("tlbp");
    index = get_cp0_index();

    // SECOND STEP: judge exception from index
    if (index & 0x80000000)
    {
        // P = 1: TLB_refill
        do_tlb_refill();
    }
    else
    {
        // P = 0: TLB_invalid
        do_tlb_invalid();
    }
}

//TOOLS:
page_frame_t *page_alloc(void)
{
    // attempt to find the first free physical frame in memory
    page_frame_t *item = (page_frame_t *)free_list.head;
    while (item != NULL)
    {
        if (item->sd_card == 0)
        {
            break;
        }
        item = item->next;
    }

    // this free frame is found
    if (item != NULL && item->sd_card == 0)
    {
        // if (clock_ptr == NULL)
        // {
        //     clock_ptr = item;
        // }

        if (item != (page_frame_t *)free_list.head)
        {
            item->prev->next = (void *)item->next;
        }
        else
        {
            free_list.head = (void *)item->next;
        }
        if (item != (page_frame_t *)free_list.tail)
        {
            item->next->prev = (void *)item->prev;
        }
        else
        {
            free_list.tail = (void *)item->prev;
        }
    }
    // cannot find a free frame in memory
    else
    {
        // choose a physical frame0 from memory (swap out, use CLOCK or sth else)
        page_frame_t *frame0;
        frame0 = swap_page_frame();

        // choose a physical frame1 from free_list which is in SD card
        page_frame_t *frame1 = (void *)free_list.head; // because you come to this step, all page_frame in free_list are in SD card
        item = frame1;
        // remove frame1 from free_list (ATTENTION! item is the head of free_list)
        if (free_list.head == free_list.tail)
        {
            free_list.head = NULL;
            free_list.tail = NULL;
        }
        else
        {
            free_list.head = (void *)(item->next);
        }
        if (item->next != NULL)
        {
            item->next->prev = NULL;
        }

        // find the adjacent page table of frame0->vpte, also swap its physical page frame out
        PTE_t *adj_pte;
        if ((frame0->vpte->vaddr & ODD_PAGE) == ODD_PAGE)
        {
            adj_pte = &(current_running[core_id]->page_table[(frame0->vpte->vaddr >> 12) - 1]);
        }
        else
        {
            adj_pte = &(current_running[core_id]->page_table[(frame0->vpte->vaddr >> 12) + 1]);
        }
        page_frame_t *frame2 = &(page_frame[adj_pte->f_pfn]);

        // choose a physical frame1 from free_list which is in SD card
        page_frame_t *frame3 = (void *)free_list.head; // because you come to this step, all page_frame in free_list are in SD card

        frame_exchange(frame0, frame1, frame2, frame3);
    }

    // add item (allocated page) to busy_list
    if (busy_list.head == NULL)
    {
        busy_list.head = (void *)item;
        busy_list.tail = (void *)item;
        item->next = NULL;
        item->prev = NULL;
    }
    else
    {
        ((page_frame_t *)busy_list.tail)->next = item;
        item->next = NULL;

        item->prev = (page_frame_t *)busy_list.tail;
        busy_list.tail = (void *)item;
    }

    return item;
}

// add paga_frame to free_list (initializing)
void page_frame_init(page_frame_t *item)
{
    if (free_list.head == NULL)
    {
        free_list.head = (void *)item;
        free_list.tail = (void *)item;
        item->next = item->prev = NULL;
    }
    else
    {
        ((page_frame_t *)free_list.tail)->next = item;
        item->next = NULL;
        item->prev = (page_frame_t *)free_list.tail;
        free_list.tail = (void *)item;
    }
}

// frame0: a choosen frame from physical memory     frame2: the PTE's mapping frame which is adjacent to frame0's PTE
// frame1: a free page frame choosen from SD card   frame3: another free page frame choosen from SD card
void frame_exchange(page_frame_t *frame0, page_frame_t *frame1, page_frame_t *frame2, page_frame_t *frame3)
{
    uint64_t temp_paddr;

    // rewrite frame0 and its adj to TLB, to assure that not trigger TLB or ADEL/ADES during memcpy and read/write_SD_card
    uint64_t entryhi;
    uint64_t entrylo0, entrylo1;
    uint32_t index;
    entryhi = (frame0->vpte->vaddr & 0xffffe000) | (frame0->asid & 0xff);
    entrylo0 = (frame0->vpte->pt_entry) >> 1;
    entrylo1 = (frame2->vpte->pt_entry) >> 1;

    set_cp0_entryhi(entryhi);
    asm volatile("tlbp");
    index = get_cp0_index();
    if (!(index & 0x80000000))
    {
        index = index;
    }
    else
    {
        index = rand_index();
    }

    set_cp0_index(index);
    set_cp0_entrylo0(entrylo0);
    set_cp0_entrylo1(entrylo1);
    asm volatile("tlbwi");

    // swap frame0 and frame1
    if (frame1 != NULL)
    {
        // use info of frame0 and frame1 to exchange data between memory and SD card
        memcpy((uint8_t *)BUFFER, (frame0->vpte->vaddr), 4096);           /* move data from frame0 (now in memory) to BUFFER */
        sd_card_read((void *)(frame0->vpte->vaddr), frame1->paddr, 4096); /* read data from frame1 (now in SD card) to frame0 */
        sd_card_write((void *)BUFFER, frame1->paddr, 4096);               /* write data from BUFFER to frame1 (this place is in SD card) */

        // rewrite frame0 and frame1 (paddr and sd_card)
        temp_paddr = frame0->paddr;
        frame0->paddr = frame1->paddr;
        frame1->paddr = temp_paddr;

        // change SD card status
        frame0->sd_card = 1;
        frame1->sd_card = 0;
        frame0->vpte->pt_entry |= 0x1;
    }

    // swap frame2 and frame3
    if ((frame2->sd_card == 0) && (frame3 != NULL))
    {
        // use info of frame2 and frame3 to exchange data between memory and SD card
        memcpy((uint8_t *)BUFFER, (frame2->vpte->vaddr), 4096);           /* move data from frame2 (now in memory) to BUFFER */
        sd_card_read((void *)(frame2->vpte->vaddr), frame3->paddr, 4096); /* read data from frame3 (now in SD card) to frame2 */
        sd_card_write((void *)BUFFER, frame3->paddr, 4096);               /* write data from BUFFER to frame3 (this place is in SD card) */

        // rewrite frame2 and frame3 (paddr and sd_card)
        temp_paddr = frame2->paddr;
        frame2->paddr = frame3->paddr;
        frame3->paddr = temp_paddr;

        // change SD card status
        frame2->sd_card = 1;
        frame3->sd_card = 0;
        frame2->vpte->pt_entry |= 0x1;
    }

    // reflush TLB: as PTE of frame0 and frame1 are adjacent, so I just need to reflush the TLB I just refill before!!
    set_cp0_entryhi(ENTRYHI_INVALID);
    set_cp0_entrylo0(0);
    set_cp0_entrylo1(0);

    asm volatile("tlbwi");
}

void get_TLB_info()
{
    uint64_t entryhi, entrylo0, entrylo1;
    int i;
    for (i = 32; i < TLB_NUM; i++)
    {
        set_cp0_index(i);
        asm volatile("tlbr");
        entryhi = get_cp0_entryhi();
        entrylo0 = get_cp0_entrylo0();
        entrylo1 = get_cp0_entrylo1();
        vt100_move_cursor(1, 20);
        if ((entrylo0 & 2))
            printk("\rTLB(%d) vpn2:%x, vaddr:%x, pfn0:%x, pfn1:%x, valid1:%d, valid2:%d, asid:%d\n", i, entryhi >> 13, entryhi & 0xffffffffffffff00, entrylo0 >> 6, entrylo1 >> 6, entrylo0 & 2, entrylo1 & 2, entryhi & 0xff);
    }
}

uint32_t rand_index()
{
    glob++;

    return (32 + glob % 32);
}

// ============================== share memory part ==========================================

// create shared memory: give a key and create a page_frame to be a shared memory
// return shmid is the subscript of page_frame array
uint64_t do_shmget(uint64_t key)
{
    int i;
    for (i = 0; i < PAGE_FRAME_NUM; i++)
    {
        if (page_frame[i].key = key)
        {
            return page_frame[i].pfn;
        }
    }

    // no current page_frame's key equal, create a new one
    page_frame_t *ptr = (page_frame_t *)free_list.head;
    while (ptr != NULL)
    {
        if (ptr->sd_card == 0)
        {
            break;
        }
        ptr = ptr->next;
    }

    if (ptr != NULL)
    {
        // remove ptr from free_list
        if (ptr == (page_frame_t *)free_list.head)
        {
            free_list.head = (void *)(ptr->next);
        }
        else
        {
            ptr->prev->next = ptr->next;
        }
        if (ptr == (page_frame_t *)free_list.tail)
        {
            free_list.tail = (void *)(ptr->prev);
        }
        else
        {
            ptr->next->prev = ptr->prev;
        }
        // add ptr to busy_list
        if (busy_list.head == NULL)
        {
            busy_list.head = (void *)ptr;
            busy_list.tail = (void *)ptr;
            ptr->next = NULL;
            ptr->prev = NULL;
        }
        else
        {
            ((page_frame_t *)busy_list.tail)->next = ptr;
            ptr->next = NULL;

            ptr->prev = (page_frame_t *)busy_list.tail;
            busy_list.tail = (void *)ptr;
        }
        // return pfn of frame
        return ptr->pfn;
    }
    else
    {
        printk("ERROR!!\n");
    }
}

// give shared memory frame a mapping: fill a blank PTE and refill TLB
uint64_t do_shmat(uint64_t shmid)
{
    int i;
    for (i = 0; i < PAGE_MAX_NUM; i++)
    {
        if ((current_running[core_id]->page_table[i].pt_entry & (PTE_V << 2)) != (PTE_V << 2))
        {
            break;
        }
    }
    current_running[core_id]->page_table[i].pt_entry = ((page_frame[shmid].paddr & 0xfffff000) >> 5 | (PTE_C << 4) | (PTE_D << 3) | (PTE_V << 2) | (page_frame[shmid].sd_card));
    current_running[core_id]->page_table[i].f_pfn = page_frame[shmid].pfn;
    uint64_t entryhi = (current_running[core_id]->page_table[i].vaddr & 0xffffffffffffe000) | (current_running[core_id]->pid & 0xff);
    uint64_t entrylo0 = (current_running[core_id]->page_table[(i / 2) * 2].pt_entry) >> 1;
    uint64_t entrylo1 = (current_running[core_id]->page_table[(i / 2) * 2 + 1].pt_entry) >> 1;

    // first search whether TLB has valid unit
    set_cp0_entryhi(entryhi);
    asm volatile("tlbp");
    uint32_t index = get_cp0_index();
    if (index & 0x80000000)
    {
        // not found, randomly fill TLB
        index = rand_index();
    }
    set_cp0_index(index);
    set_cp0_entryhi(entryhi);
    set_cp0_entrylo0(entrylo0);
    set_cp0_entrylo1(entrylo1);
    asm volatile("tlbwi");
    if (i == 0)
    {
        return current_running[core_id]->page_table[i].vaddr + 0x8;
    }
    else
    {
        return current_running[core_id]->page_table[i].vaddr;
    }
}

// reflush TLB
uint64_t do_shmdt(uint64_t addr)
{
    current_running[core_id]->page_table[addr >> 12].pt_entry &= 0xfffffffb; // V = 0
    uint64_t entryhi = (addr & 0xffffffffffffe000) | (current_running[core_id]->pid & 0xff);
    set_cp0_entryhi(entryhi);
    asm volatile("tlbp");
    uint32_t index = get_cp0_index();
    if (!(index & 0x80000000))
    {
        set_cp0_entryhi(ENTRYHI_INVALID);
        set_cp0_entrylo0(0);
        set_cp0_entrylo1(0);
        asm volatile("tlbwi");
    }
}